package edu.osu.intelligentagents.common;

import java.util.HashMap;

public class NonStationaryPolicy extends Policy {
	private HashMap<String, Action> policies = new HashMap<>();
	private MDP mdp;

	public NonStationaryPolicy() {
	}

	public void setMDP(MDP mdp) {
		this.mdp = mdp;

	}

	// Bellman backup and value iteration.
	public void valueIteration(float horizon) {
		for (int h = 1; h < horizon + 1; h++) {
			for (State state : mdp.states) {
				float maxValue = -Float.MAX_VALUE;
				for (Action action : mdp.actions) {
					float value = 0;
					for (State state_prime : mdp.states) {
						value += action.getStates(state.getId()).get(
								state_prime.getId())
								* state_prime.getValues().get(h - 1);
					}
					if (value > maxValue) {
						maxValue = value;
						this.setAction(state, h, action);
					}
				}
				state.getValues().add(maxValue + state.getReward());
			}
		}

	}

	private static String getKey(State state, int timeRemain) {
		return "s=" + state.getId() + ",k=" + timeRemain;
	}

	@Override
	public Action getAction(State state) {
		throw new RuntimeException(
				"This method can not be called on a Non StationaryPolicy Policy");
	}

	@Override
	public Action getAction(State state, int timeRemain) {
		// TODO Auto-generated method stub
		Action action = policies.get(getKey(state, timeRemain));
		if (action == null)
			return new Action(-1, "NOP", null);
		return action;
	}

	public void setAction(State state, int timeRemian, Action action) {
		policies.put(getKey(state, timeRemian), action);
	}
}
